import os
import unittest
import numpy as np
import pandas as pd
from api.base import CommonFunction
from api.preprocessing.NullHandler import NullHandler


class NullHandlerTestCase(unittest.TestCase):
    def _reset_data(self):
        self.data = pd.DataFrame([[1, 2, 3, 4, 5],
                                  [2, None, 4, 5, 6],
                                  [3, 4, None, 6, 7],
                                  [4, 5, 6, 7, 8]], columns=list('ABCDE'))
        self.null_handler = NullHandler(data=self.data)

    def _judge_same_dataframe(self, target_dataframe):
        # 判断空值位置是否相同
        self.assertTrue((np.isnan(self.null_handler.get_data().values) == np.isnan(target_dataframe.values)).all())
        # 判断所有值是否相同
        self.assertTrue((self.null_handler.get_data().fillna(value=0).values == target_dataframe.fillna(value=0).values).all())

    def test_init(self):
        self._reset_data()
        target_dataframe = pd.DataFrame([[1, 2, 3, 4, 5],
                                         [2, None, 4, 5, 6],
                                         [3, 4, None, 6, 7],
                                         [4, 5, 6, 7, 8]], columns=list('ABCDE'))
        self._judge_same_dataframe(target_dataframe)

        self.data = pd.DataFrame([[1, 2, 3, 4, 5],
                                  [2, None, 4, 5, 6],
                                  [3, 4, None, 6, 7],
                                  [4, 5, 6, 7, 8]], columns=list('ABCDE'))
        self.null_handler = NullHandler(data=self.data, columns=list('ABCD'))
        target_dataframe = pd.DataFrame([[1, 2, 3, 4],
                                         [2, None, 4, 5],
                                         [3, 4, None, 6],
                                         [4, 5, 6, 7]], columns=list('ABCD'))
        self._judge_same_dataframe(target_dataframe)

    def test_dropna_axis(self):
        self._reset_data()
        self.null_handler.dropna(axis=0)
        target_dataframe = pd.DataFrame([[1, 2, 3, 4, 5],
                                         [4, 5, 6, 7, 8]], columns=list('ABCDE'))
        self._judge_same_dataframe(target_dataframe)

        self._reset_data()
        self.null_handler.dropna(axis=1)
        target_dataframe = pd.DataFrame([[1, 4, 5],
                                         [2, 5, 6],
                                         [3, 6, 7],
                                         [4, 7, 8]], columns=list('ADE'))
        self._judge_same_dataframe(target_dataframe)

    def test_dropna_null_value(self):
        self._reset_data()
        self.null_handler.dropna(null_value=1)
        target_dataframe = pd.DataFrame([[2, None, 4, 5, 6],
                                         [3, 4, None, 6, 7],
                                         [4, 5, 6, 7, 8]], columns=list('ABCDE'))
        self._judge_same_dataframe(target_dataframe)

        self._reset_data()
        self.null_handler.dropna(null_value=1, axis=1)
        target_dataframe = pd.DataFrame([[2, 3, 4, 5],
                                         [None, 4, 5, 6],
                                         [4, None, 6, 7],
                                         [5, 6, 7, 8]], columns=list('BCDE'))
        self._judge_same_dataframe(target_dataframe)

    def test_fillna_mean(self):
        self._reset_data()
        self.null_handler.fillna(axis=0, method='mean')
        target_dataframe = pd.DataFrame([[1, 2, 3, 4, 5],
                                         [2, 4.25, 4, 5, 6],
                                         [3, 4, 5, 6, 7],
                                         [4, 5, 6, 7, 8]], columns=list('ABCDE'))
        self._judge_same_dataframe(target_dataframe)

    def test_fillna_median(self):
        self._reset_data()
        self.null_handler.fillna(axis=0, method='median')
        target_dataframe = pd.DataFrame([[1, 2, 3, 4, 5],
                                         [2, 4.5, 4, 5, 6],
                                         [3, 4, 5, 6, 7],
                                         [4, 5, 6, 7, 8]], columns=list('ABCDE'))
        self._judge_same_dataframe(target_dataframe)

    def test_fillna_ffill(self):
        self._reset_data()
        self.null_handler.fillna(axis=0, method='ffill')
        target_dataframe = pd.DataFrame([[1, 2, 3, 4, 5],
                                         [2, 2, 4, 5, 6],
                                         [3, 4, 4, 6, 7],
                                         [4, 5, 6, 7, 8]], columns=list('ABCDE'))
        self._judge_same_dataframe(target_dataframe)

        self._reset_data()
        self.null_handler.fillna(axis=1, method='ffill')
        target_dataframe = pd.DataFrame([[1, 2, 3, 4, 5],
                                         [2, 2, 4, 5, 6],
                                         [3, 4, 4, 6, 7],
                                         [4, 5, 6, 7, 8]], columns=list('ABCDE'))
        self._judge_same_dataframe(target_dataframe)

    def test_fillna_bfill(self):
        self._reset_data()
        self.null_handler.fillna(axis=0, method='bfill')
        target_dataframe = pd.DataFrame([[1, 2, 3, 4, 5],
                                         [2, 4, 4, 5, 6],
                                         [3, 4, 6, 6, 7],
                                         [4, 5, 6, 7, 8]], columns=list('ABCDE'))
        self._judge_same_dataframe(target_dataframe)

        self._reset_data()
        self.null_handler.fillna(axis=1, method='bfill')
        target_dataframe = pd.DataFrame([[1, 2, 3, 4, 5],
                                         [2, 4, 4, 5, 6],
                                         [3, 4, 6, 6, 7],
                                         [4, 5, 6, 7, 8]], columns=list('ABCDE'))
        self._judge_same_dataframe(target_dataframe)

    def test_fillna_const(self):
        self._reset_data()
        self.null_handler.fillna(axis=0, method='const', value=-9999)
        target_dataframe = pd.DataFrame([[1, 2, 3, 4, 5],
                                         [2, -9999, 4, 5, 6],
                                         [3, 4, -9999, 6, 7],
                                         [4, 5, 6, 7, 8]], columns=list('ABCDE'))
        self._judge_same_dataframe(target_dataframe)


if __name__ == '__main__':
    unittest.main()


